#include <cutil_inline.h>
#include <cutil_math.h>
#include <math_constants.h>
#include "radixsort.h"
#include "CudaUtils.h"
#include "PMHash.h"
#include "PMHash_Kernels.cu"

struct HashMap
{
	DevicePhotons m_DevPhotons;
	HashTable m_HashTable;
};

DevicePhotons AllocHashDevicePhotons(unsigned int p_Elements)
{
	DevicePhotons photons;
	cutilSafeCall(CudaMallocRec((void**)(&photons.m_Position), sizeof(float4) * p_Elements));
	cutilSafeCall(CudaMallocRec((void**)(&photons.m_Power), sizeof(uchar4) * p_Elements));
	return photons;
}

void FreeHashDevicePhotons(DevicePhotons p_Photons)
{
	cutilSafeCall(CudaFreeRec(p_Photons.m_Position));
	cutilSafeCall(CudaFreeRec(p_Photons.m_Power));
}

HashPositions AllocHashPositions(unsigned int p_Elements)
{
	HashPositions hashPos;
	cutilSafeCall(CudaMallocRec((void**)(&hashPos.m_HashValue), p_Elements * sizeof(int)));
	cutilSafeCall(CudaMallocRec((void**)(&hashPos.m_PhotonIdx), p_Elements * sizeof(uint)));
	return hashPos;
}

void FreeHashPositions(HashPositions p_HashPos)
{
	cutilSafeCall(CudaFreeRec(p_HashPos.m_HashValue));
	cutilSafeCall(CudaFreeRec(p_HashPos.m_PhotonIdx));
}

HashTable AllocHashTable(unsigned int p_Elements)
{
	HashTable ht;
	cutilSafeCall(CudaMallocRec((void**)(&ht.m_FirstIdx), p_Elements * sizeof(int)));
	cutilSafeCall(CudaMallocRec((void**)(&ht.m_NumPhotons), p_Elements * sizeof(uint)));
	return ht;
}

void FreeHashTable(HashTable p_HT)
{
	cutilSafeCall(CudaFreeRec(p_HT.m_FirstIdx));
	cutilSafeCall(CudaFreeRec(p_HT.m_NumPhotons));
}


void* PerformHashing(void* p_Photons, unsigned int p_Elements, float* p_BBMin, float* p_BBMax, 
					 float p_MaxRadius, unsigned int p_Timer)
{
	DevicePhotons devPhotons = AllocHashDevicePhotons(p_Elements);

	// set hash parameters
	float3 bbMin = *((float3*)p_BBMin);
	float3 bbMax = *((float3*)p_BBMax);
	float cellSize = p_MaxRadius;
	HashParams hashParams;
	hashParams.m_WorldOrigin.x = bbMin.x - (2.0f * cellSize);
	hashParams.m_WorldOrigin.y = bbMin.y - (2.0f * cellSize);
	hashParams.m_WorldOrigin.z = bbMin.z - (2.0f * cellSize);
	hashParams.m_CellSize = cellSize;
	hashParams.m_CellScaleFactor.x = hashParams.m_CellScaleFactor.y = hashParams.m_CellScaleFactor.z = 1.0f / cellSize;
	hashParams.m_GridSize.x = (uint)(((bbMax.x - bbMin.x) + (4.0f * cellSize)) / cellSize + 1.0f);
	hashParams.m_GridSize.y = (uint)(((bbMax.y - bbMin.y) + (4.0f * cellSize)) / cellSize + 1.0f);
	hashParams.m_GridSize.z = (uint)(((bbMax.z - bbMin.z) + (4.0f * cellSize)) / cellSize + 1.0f);
	cudaMemcpyToSymbol(g_Params, &hashParams, sizeof(HashParams));

	// copy input data on the device
	HostPhoton* hostPhotons;
	cutilSafeCall(CudaMallocRec((void**)(&hostPhotons), sizeof(HostPhoton) * p_Elements));
	cutilSafeCall(cudaMemcpy(hostPhotons, p_Photons, sizeof(HostPhoton) * p_Elements, cudaMemcpyHostToDevice));

	// convert input data
	int numBlocks, numThreads;
	ComputeGridSize(p_Elements, 256, numBlocks, numThreads);
	ConvertInputData<<<numBlocks, numThreads>>>(hostPhotons, devPhotons, p_Elements);
	CudaCheckError("ConvertInputData execution failed");
	cutilSafeCall(CudaFreeRec(hostPhotons));

	// start construction timer
	CUT_SAFE_CALL(cutStartTimer(p_Timer));

	// generate offset lookup table
	int curEntry = 0;
	float4 offsetTable[27];
	for(int i = 0; i < 3; ++i)
	{
		for(int j = 0; j < 3; ++j)
		{
			for(int k = 0; k < 3; ++k)
			{
				offsetTable[curEntry].x = (i * cellSize) - cellSize;
				offsetTable[curEntry].y = (j * cellSize) - cellSize;
				offsetTable[curEntry].z = (k * cellSize) - cellSize;
				++curEntry;
			}
		}
	}

	// set the middle cell (0/0/0 offset) as the first
	offsetTable[13] = offsetTable[0];
	offsetTable[0].x = offsetTable[0].y = offsetTable[0].z = 0.0f;
	cutilSafeCall(cudaMemcpyToSymbol(g_CellOffsets, offsetTable, sizeof(float4) * 27));

	// create hash position list
	HashPositions hashPositions = AllocHashPositions(p_Elements);
 	CalcPositionHashes<<<numBlocks, numThreads>>>(devPhotons.m_Position, hashPositions, p_Elements);
 	CudaCheckError("CalcHashPositions execution failed");

 	// sort hash position list
 	RadixSort rs(p_Elements);
 	rs.sort(hashPositions.m_HashValue, hashPositions.m_PhotonIdx, p_Elements, 32);
 
 	// reorder photons
	DevicePhotons buffer = AllocHashDevicePhotons(p_Elements);
 	ReorderPhotons<true><<<numBlocks, numThreads>>>(devPhotons, buffer, hashPositions.m_PhotonIdx, p_Elements);
 	CudaCheckError("ReorderPhotons<true> execution failed");
	ReorderPhotons<false><<<numBlocks, numThreads>>>(devPhotons, buffer, hashPositions.m_PhotonIdx, p_Elements);
	CudaCheckError("ReorderPhotons<false> execution failed");
 
 	// create hash lookup table
#if defined USE_SPATIAL_HASHING
	const uint numHashes = SpatialHashTableSize;
#else
	const uint numHashes = hashParams.m_GridSize.x * hashParams.m_GridSize.y * hashParams.m_GridSize.z;
#endif
 	HashTable hashTable = AllocHashTable(numHashes);
 	//ComputeGridSize(numHashes, 256, numBlocks, numThreads);
 	CreateHashTable<<<numBlocks, numThreads>>>(hashPositions.m_HashValue, hashTable, p_Elements, numHashes);
 	CudaCheckError("CreateHashTable execution failed");

	// free unneeded structures
	FreeHashPositions(hashPositions);

	// return hash map handle
	HashMap* hm = (HashMap*)malloc(sizeof(HashMap));
	hm->m_DevPhotons = devPhotons;
	hm->m_HashTable = hashTable;

	// stop construction timer
	CUT_SAFE_CALL(cutStopTimer(p_Timer));
	return (void*)hm;
}

// void HashSearch(void* p_HM, float* p_Queries, unsigned int p_NumPoints, float p_MaxDist, unsigned int p_MaxPhotons, void* p_SqDist, void* p_Powers)
// {
// 	HashMap* hashMap = (HashMap*)p_HM;
// 
// 	// allocate memory for the query positions and the results
// 	float4* query;
// 	cutilSafeCall(cudaMalloc((void**)&query, sizeof(float4) * p_NumPoints));
// 	cutilSafeCall(cudaMemcpy(query, p_Queries, sizeof(float4) * p_NumPoints, cudaMemcpyHostToDevice));
// 
// 	float* tmpSqDist;
// 	uchar4* tmpPowers;
// 	cutilSafeCall(cudaMalloc((void**)&tmpSqDist, sizeof(float) * p_NumPoints * p_MaxPhotons * 27));
// 	cutilSafeCall(cudaMalloc((void**)&tmpPowers, sizeof(uchar4) * p_NumPoints * p_MaxPhotons * 27));
// 
// 	float* sqDist;
// 	uchar4* powers;
// 	cutilSafeCall(cudaMalloc((void**)&sqDist, sizeof(float) * p_NumPoints * p_MaxPhotons));
// 	cutilSafeCall(cudaMalloc((void**)&powers, sizeof(uchar4) * p_NumPoints * p_MaxPhotons));
// 
// 	// find photons
// 	dim3 t_3x3x3(3, 3, 3);
// 	dim3 b_512(1, 1, 1);
// 	if(p_NumPoints > 65535)
// 	{
// 		b_512.x = 2048;
// 		b_512.y = p_NumPoints / 2048;
// 		if(p_NumPoints % 2048) ++b_512.y;
// 	}
// 	else
// 		b_512.x = p_NumPoints;
// 
// 	const uint SharedMem = (sizeof(float) + sizeof(uchar4)) * 27 * p_MaxPhotons;
// 	PhotonSearch<<<b_512, t_3x3x3, SharedMem>>>
// 		(hashMap->m_DevPhotons, hashMap->m_HashTable, query, p_NumPoints, 
// 		 p_MaxDist * p_MaxDist, p_MaxPhotons, tmpSqDist, tmpPowers);
// 	CudaCheckError("PhotonSearch execution failed");
// 
// 	int numThreads, numBlocks;
// 	ComputeGridSize(p_NumPoints, 256, numBlocks, numThreads);
// 	FinalizeResults<<<numBlocks, numThreads>>>(tmpSqDist, tmpPowers, sqDist, powers, p_MaxPhotons, p_NumPoints);
// 	CudaCheckError("FinalizeResults execution failed");
// 
// 	// copy results back to host memory
// 	cutilSafeCall(cudaMemcpy(p_SqDist, sqDist, sizeof(float) * p_NumPoints * p_MaxPhotons, cudaMemcpyDeviceToHost));
// 	cutilSafeCall(cudaMemcpy(p_Powers, powers, sizeof(uchar4) * p_NumPoints * p_MaxPhotons, cudaMemcpyDeviceToHost));
// 
// 	// free memory
// 	cutilSafeCall(cudaFree(tmpSqDist));
// 	cutilSafeCall(cudaFree(tmpPowers));
// 	cutilSafeCall(cudaFree(sqDist));
// 	cutilSafeCall(cudaFree(powers));
// 	cutilSafeCall(cudaFree(query));
// }

void HashSearch(void* p_HM, float* p_Queries, unsigned int p_NumPoints, float p_MaxDist, 
				unsigned int p_MaxPhotons, void* p_SqDist, void* p_Powers)
{
	HashMap* hashMap = (HashMap*)p_HM;

	// allocate memory for the query positions and the results
	float4* query;
	cutilSafeCall(cudaMalloc((void**)&query, sizeof(float4) * p_NumPoints));
	cutilSafeCall(cudaMemcpy(query, p_Queries, sizeof(float4) * p_NumPoints, cudaMemcpyHostToDevice));

	float* sqDist;
	uchar4* powers;
	cutilSafeCall(cudaMalloc((void**)&sqDist, sizeof(float) * p_NumPoints * p_MaxPhotons));
	cutilSafeCall(cudaMalloc((void**)&powers, sizeof(uchar4) * p_NumPoints * p_MaxPhotons));

	// find photons
	int numBlocks, numThreads;
	const uint SharedMem = (sizeof(float) + sizeof(uchar4)) * p_MaxPhotons * 32; // 32 Warp size
	ComputeGridSize(p_NumPoints, 32, numBlocks, numThreads);
	PhotonSearch2<<<numBlocks, numThreads, SharedMem>>>
		(hashMap->m_DevPhotons, hashMap->m_HashTable, query, p_NumPoints, 
		p_MaxDist * p_MaxDist, p_MaxPhotons, sqDist, powers);
	CudaCheckError("PhotonSearch2 execution failed");

	// copy results back to host memory
	cutilSafeCall(cudaMemcpy(p_SqDist, sqDist, sizeof(float) * p_NumPoints * p_MaxPhotons, cudaMemcpyDeviceToHost));
	cutilSafeCall(cudaMemcpy(p_Powers, powers, sizeof(uchar4) * p_NumPoints * p_MaxPhotons, cudaMemcpyDeviceToHost));

	// free memory
	cutilSafeCall(cudaFree(sqDist));
	cutilSafeCall(cudaFree(powers));
	cutilSafeCall(cudaFree(query));
}

void CleanUpHash(void* p_HM)
{
	HashMap* hashMap = (HashMap*)p_HM;
	FreeHashTable(hashMap->m_HashTable);
	FreeHashDevicePhotons(hashMap->m_DevPhotons);
	free(hashMap);
}